package com.biubiu.doubaoai;

import android.annotation.SuppressLint;
import android.util.Log;

import com.biubiu.doubaoai.listener.CallBackListener;
import com.volcengine.ark.runtime.exception.ArkHttpException;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatFunction;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.model.completion.chat.ChatTool;
import com.volcengine.ark.runtime.service.ArkService;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import io.reactivex.Flowable;
import io.reactivex.android.schedulers.AndroidSchedulers;
import io.reactivex.schedulers.Schedulers;


public class ChatCompletions {
    public static void withChatAi(CallBackListener listener) {

        ArkService service = ArkService.builder().apiKey("d1777345-3843-47ac-8884-8b472b9b2b32").build();

        System.out.println("\n----- standard request -----");
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("写出《登鹳鹊楼》全诗，只需简洁回复。").build();
        messages.add(systemMessage);
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20241030112518-k7757")
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
            String text = choice.getMessage().getContent().toString();
            Log.d("liaoyucheng", text);
            listener.callback(text);
        });

        // shutdown service
        service.shutdownExecutor();
    }

    public static void withChatAi2(CallBackListener listener) {

        ArkService service = ArkService.builder().apiKey("d1777345-3843-47ac-8884-8b472b9b2b32").build();

        System.out.println("\n----- multiple rounds request -----");
        final List<ChatMessage> messages = Arrays.asList(
                ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build(),
                ChatMessage.builder().role(ChatMessageRole.USER).content("花椰菜是什么？").build(),
                ChatMessage.builder().role(ChatMessageRole.ASSISTANT).content("花椰菜又称菜花、花菜，是一种常见的蔬菜。").build(),
                ChatMessage.builder().role(ChatMessageRole.USER).content("再详细点").build()
        );

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20241030112518-k7757")
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
            String text = choice.getMessage().getContent().toString();
            Log.d("liaoyucheng", text);
            listener.callback(text);
        });

        // shutdown service
        service.shutdownExecutor();
    }

    /**
     * 流试-同步
     */
    @SuppressLint("CheckResult")
    public static void withChatAi3(CallBackListener listener) {

        ArkService service = ArkService.builder().apiKey("d1777345-3843-47ac-8884-8b472b9b2b32").build();

        System.out.println("\n----- streaming request -----");
        final List<ChatMessage> streamMessages = new ArrayList<>();
        final ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
        final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("写出3首李白的诗，只需简洁回复。").build();
        streamMessages.add(streamSystemMessage);
        streamMessages.add(streamUserMessage);

        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20241030112518-k7757")
                .messages(streamMessages)
                .build();

        try {
            long time1 = System.currentTimeMillis();
            Flowable<ChatCompletionChunk> flowable = service.streamChatCompletion(streamChatCompletionRequest);
            long time2 = System.currentTimeMillis();
            Log.d("liaoyucheng", "耗时："+(time2-time1));
            flowable.doOnError(Throwable::printStackTrace)
                    .subscribeOn(Schedulers.io())
                    .observeOn(Schedulers.single())
                    .doOnComplete(
                            ()->{
                                service.shutdownExecutor();
                                Log.d("liaoyucheng", "结束");
                            }
                    )
                    .forEach(
                            choice -> {
                                if (!choice.getChoices().isEmpty()) {
                                    String text = choice.getChoices().get(0).getMessage().getContent().toString();
                                    Log.d("liaoyucheng", text);
                                    listener.callback(text);
                                }
                            }
                    );
        } catch (ArkHttpException e) {
            System.out.print(e.toString());
        }
        Log.d("liaoyucheng", "阻塞结束");
        // shutdown service
        //service.shutdownExecutor();
    }

    /**
     * 流试-异步
     */
    @SuppressLint("CheckResult")
    public static void withChatAi4(CallBackListener listener) {

        ArkService service = ArkService.builder().apiKey("d1777345-3843-47ac-8884-8b472b9b2b32").build();

        System.out.println("\n----- streaming request -----");
        final List<ChatMessage> streamMessages = new ArrayList<>();
        final ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
        final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("写出3首李白的诗，只需简洁回复。").build();
        streamMessages.add(streamSystemMessage);
        streamMessages.add(streamUserMessage);

        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20241030112518-k7757")
                .messages(streamMessages)
                .build();

        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .subscribeOn(Schedulers.io())
                .observeOn(Schedulers.single())
                .subscribe(
                        choice -> {
                            if (!choice.getChoices().isEmpty()) {
                                String text = choice.getChoices().get(0).getMessage().getContent().toString();
                                Log.d("liaoyucheng", text);
                                listener.callback(text);
                            }
                        },
                        error -> {
                            error.printStackTrace();
                            Log.d("liaoyucheng", "onError");
                        },
                        () -> {
                            service.shutdownExecutor();
                            Log.d("liaoyucheng", "onComplete");
                        }
                );
    }

    public static void withChatAi5(CallBackListener listener) {
        ArkService service = ArkService.builder().apiKey("d1777345-3843-47ac-8884-8b472b9b2b32").build();

        System.out.println("\n----- function call request -----");
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("北京今天天气如何？").build();
        messages.add(userMessage);

        final List<ChatTool> tools = Arrays.asList(
                new ChatTool(
                        "function",
                        new ChatFunction.Builder()
                                .name("get_current_weather")
                                .description("获取给定地点的天气")
                                .parameters(new Weather(
                                        "object",
                                        new HashMap<String, Object>() {{
                                            put("location", new HashMap<String, String>() {{
                                                put("type", "string");
                                                put("description", "T地点的位置信息，比如北京");
                                            }});
                                            put("unit", new HashMap<String, Object>() {{
                                                put("type", "string");
                                                put("enum", Arrays.asList("摄氏度", "华氏度"));
                                            }});
                                        }},
                                        Collections.singletonList("location")
                                ))
                                .build()
                )
        );


        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20241030112518-k7757")
                .messages(messages)
                .tools(tools)
                .build();

        //service.createChatCompletion(chatCompletionRequest).getChoices().forEach(System.out::println);
        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
            String text = choice.getMessage().getContent().toString();
            Log.d("liaoyucheng", text);
            listener.callback(text);
        });

        // shutdown service
        service.shutdownExecutor();
    }

    public static class Weather {
        public String type;
        public Map<String, Object> properties;
        public List<String> required;

        public Weather(String type, Map<String, Object> properties, List<String> required) {
            this.type = type;
            this.properties = properties;
            this.required = required;
        }

        public String getType() {
            return type;
        }

        public void setType(String type) {
            this.type = type;
        }

        public Map<String, Object> getProperties() {
            return properties;
        }

        public void setProperties(Map<String, Object> properties) {
            this.properties = properties;
        }

        public List<String> getRequired() {
            return required;
        }

        public void setRequired(List<String> required) {
            this.required = required;
        }
    }
}